#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Copyright 2020 Huawei Technologies Co., Ltd
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import sys
import struct
import acl
import numpy as np
import cv2
sys.path.append("..")
from base.dvpp import DvppInfo
from common import const
from common.log import Log, check_ret
from resource.resource import ascend_resource


class PostProcess():
    def __init__(self, context, 
                model_input_height, model_input_width,
                image_height, image_width, thread_arrange):
        self.class_name = self.__class__.__name__
        self.context = context
        self.model_input_height = model_input_height
        self.model_input_width = model_input_width
        self.image_height = image_height
        self.image_width = image_width
        self.thread_arrange = thread_arrange
        self.memory_pool = ascend_resource.memory_pool
        self.dst_host = None
        self.buffer_size = 0
        self.ratio = 0
        self.create_post_buffer()
    
    def create_post_buffer(self):
        # malloc for each batch of reference images
        width_stride, height_stride = \
            self.memory_pool.stride_calculation(self.image_height,
                                            self.image_width,
                                            'vdec')
        self.buffer_size = \
            self.memory_pool.buffer_size_calculation(width_stride,
                                height_stride,
                                const.PIXEL_FORMAT_YUV_SEMIPLANAR_420)
        self.dst_host = \
            self.memory_pool.memory_allocation_(self.buffer_size * \
                                                self.thread_arrange, 
                                                const.HOST)
        self.ratio =  self.image_width / self.model_input_width

    def send_to_post(self, yuv_batch_ref):
        for j in range(0, self.thread_arrange):
            ret = acl.rt.memcpy(self.dst_host + j * self.buffer_size, 
                                yuv_batch_ref[j][1],
                                yuv_batch_ref[j][0], 
                                yuv_batch_ref[j][1],
                                const.ACL_MEMCPY_DEVICE_TO_HOST)
            check_ret(self.class_name, 'postprocess',
                    'fail to memcpy to host', ret)
            self.memory_pool.release_memory(yuv_batch_ref[j][0], const.DVPP)

    def _unpack_output_item(self, byte_array, shape, datatype):
        tag = ""
        np_type = None
        if datatype == const.ACL_FLOAT:
            tag = 'f'
            np_type = np.float
        elif datatype == const.ACL_INT32:
            tag = 'i'
            np_type = np.int32
        elif datatype == cont.ACL_UINT32:
            tag = 'I'
            np_type = np.uint32
        else:
            check_ret(self.class_name, '_unpack_output_item', 
                'unsurpport datatype', 1)
            return
        size = byte_array.size / 4
        unpack_tag = str(int(size)) + tag
        data = struct.unpack(unpack_tag, bytearray(byte_array))
        return np.array(data).astype(np_type).reshape(shape)
    
    def run_postprocess(self, frame_id, channel_id, 
                    infer_result, yuv_batch_ref, output_info):
        # send to post
        self.send_to_post(yuv_batch_ref)
        # bounding box
        box_info = acl.util.ptr_to_numpy(infer_result[0][0], 
                                        (infer_result[0][1],), 
                                        const.NPY_BYTE)
        box_num = acl.util.ptr_to_numpy(infer_result[1][0], 
                                        (infer_result[1][1],), 
                                        const.NPY_BYTE)
        box_info = self._unpack_output_item(box_info, 
                                            output_info[0]["shape"], 
                                            output_info[0]["type"])
        box_num = self._unpack_output_item(box_num, 
                                        output_info[1]["shape"], 
                                        output_info[1]["type"])
        self.memory_pool.release_memory(infer_result[0][0], const.HOST)
        self.memory_pool.release_memory(infer_result[1][0], const.HOST)
        # post process
        # yolo v3
        # { TOPLEFTX = 0, TOPLEFTY, BOTTOMRIGHTX, BOTTOMRIGHTY, SCORE, LABEL }
        for j in range(0, self.thread_arrange):
            total_box = box_num[j, 0]
            host_ptr = self.dst_host + j * self.buffer_size
            out = acl.util.ptr_to_numpy(host_ptr, 
                                        (yuv_batch_ref[j][3] * 3 // 2, 
                                        yuv_batch_ref[j][2]), 
                                        2)
            img_rgb = cv2.cvtColor(out, cv2.COLOR_YUV2RGB_NV21)
            for i in range(0, total_box):
                score = int((box_info[j, const.SCORE * total_box + i]) * 100)
                top_x = int(box_info[j, i + total_box * const.TOPLEFTX] * \
                            self.ratio)
                top_y = int(box_info[j, i + total_box * const.TOPLEFTY] * \
                            self.ratio)
                bottom_x = int(box_info[j, i + total_box * const.BOTTOMRIGHTX] * \
                            self.ratio)
                bottom_y = int(box_info[j, i + total_box * const.BOTTOMRIGHTY] * \
                            self.ratio)
                label = int(box_info[j, i + total_box * const.LABEL])
                text = const.yolov3Label[label] + str(score) + '%'
                print('top_xy:{} {} bottom_xy:{} {} label:{} text:{}'.\
                    format(top_x, top_y, bottom_x, bottom_y, label, text))
                cv2.rectangle(img_rgb, (top_x, top_y), 
                            (bottom_x, bottom_y), (255, 0, 0), 2)
                cv2.putText(img_rgb, text, (top_x, top_y-2), 
                            cv2.FONT_HERSHEY_COMPLEX, 1.5, (0, 0, 255), 2)
            print('----------channel_id: {}   frame_id: {}-------------\n'.\
                    format(channel_id[j], frame_id[j])) 
            #cv2.imwrite('./image_' + str(index) + '.jpg', img_rgb)
            #img_resize = cv2.resize(img_rgb, 
            #    (self.image_width // 4, self.image_height // 4))
            #cv2.imshow('result', img_resize)
            #cv2.waitKey(1)
    
    def release_postprocess(self):
        ret = acl.rt.set_context(self.context)
        #cv2.destroyAllWindows()
        if self.dst_host:
            self.memory_pool.release_memory(self.dst_host, const.HOST)
